package com.taoge.model;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.LongAdder;

/**
 * 漏桶限流类
 * @author Mr.Tao
 */
public final class LeakyBucket {

    /**
     * 漏桶集合
     */
    private static final Map<String, LeakyBucket> LEAKY_BUCKET_MAP = new ConcurrentHashMap<>();

    /**
     * 桶的容量
     */
    private final int capacity;

    /**
     * 木桶剩余的水滴的量(初始化的时候的空的桶)
     */
    private final LongAdder water = new LongAdder();

    /**
     * 水滴的流出的速率   leakRate==1000相当于每1000毫秒流出1滴
     */
    private final int leakRate;

    /**
     * 第一次请求之后,木桶在这个时间点开始漏水
     */
    private long leakTimeStamp;

    public LeakyBucket(int capacity, int leakRate) {
        this.capacity = capacity;
        this.leakRate = leakRate;
    }

    public boolean isFull() {
        // 如果是空桶，就当前时间作为桶开是漏出的时间
        long currentTime = System.currentTimeMillis();
        if (water.sum() == 0) {
            leakTimeStamp = currentTime;
            water.increment();
            return false;
        }

        // 先执行漏水，计算剩余水量
        long waterLeft = water.sum() - ((currentTime - leakTimeStamp) / leakRate);
        water.reset();
        if (waterLeft > 0) {
            water.add(waterLeft);
        }

        // 重新更新leakTimeStamp
        leakTimeStamp = currentTime;

        // 尝试加水,并且水还未满
        if (water.sum() < capacity) {
            water.increment();
            return false;
        }

        // 水满，拒绝加水
        return true;
    }

    public static void putLeakyBucket(String key, int capacity, int leakRate) {
        LEAKY_BUCKET_MAP.put(key, new LeakyBucket(capacity, leakRate));
    }

    public static LeakyBucket getLeakyBucket(String key) {
        return LEAKY_BUCKET_MAP.get(key);
    }

}
